import os
from typing import Any
from .utils import load_csv, load_txt


class Graph:
    def __init__(self, graph_dir="data/Graph"):

        self.id2ent = load_txt(os.path.join(graph_dir, "id2ent.txt"))
        self.ent2id = {ent: idx for idx, ent in enumerate(self.id2ent)}
        self.id2rel = load_txt(os.path.join(graph_dir, "id2rel.txt"))
        self.rel2id = {rel: idx for idx, rel in enumerate(self.id2rel)}

        self.all_entids = list(self.ent2id.values())
        self.all_relids = list(self.rel2id.values())

        _, triples = load_csv(os.path.join(graph_dir, "graph.csv"))

        self.triples_list = [tuple(triple) for triple in triples]
        self.triples_set = set(self.triples_list)

        return None

    def __getattribute__(self, __name: str) -> Any:

        if __name == "num_entities":
            return len(self.id2ent)
        elif __name == "num_relations":
            return len(self.id2rel)
        else:
            return super().__getattribute__(__name)

    def __len__(self):
        return len(self.triples_list)

    def __getitem__(self, idx):
        return self.triples_list[idx]

    def __contains__(self, triple):
        return triple in self.triples_set

    def __iter__(self):
        return iter(self.triples_list)

    def __repr__(self):
        return f"{self.__class__.__name__}(len={len(self)}, triples={self.triples_list[:3]})"
